vllm-ascend PR #18870 解读：FlashComm1 的"假开启"与 SP 归位

1. PR 档案

2. 问题：guard 一刀切，FlashComm1 名开实关（改动前）

先理清四个角色怎么协作。TP=16 把一层切到 16 卡；SP 在 TP 组内再把 token 横切成连续段（每卡一段）；PP 再把模型的 N 层按层切成分段（stage）接力——层是 PP 的切法，token 是 SP 的切法，两者正交。GLM DSA V1 开全局 SP（FlashComm1 模式）时，decode 侧的分工是：注意力投影不自带归约，由 decoder 收口统一归约——这样每层省一次集合通信，这正是 FlashComm1 的收益来源。

但上游 vLLM 立了一条 guard：PP>1 时关掉 decoder 级 SP。原因在上游自己的传输约定——它的 PP transport 期望拿到全量激活，如果 stage 内跑 SP、出口传出的是分片，上游的接法就对不上。guard 一刀切之后，PP>1 的 GLM V1 就出现了下面的名实不符（PP 切层与 SP 的关系见图 3）：

ⒶⒷ 速查（后文正文与代码注释均用此编号）

3. 核心思路：自己兑现上游的前提

guard 的担忧是对的——上游 transport 确实只认全量。本 PR 不跟 guard 对着干，而是把 guard 的前提自己兑现：stage 内部照开 SP（拿 FlashComm1 的收益），在每个 PP 出口把 hidden/residual 聚回全量再交给 transport；o_proj 的归约权从"每层各归各的"收归 decoder 统一收口；Runner 侧所有 PP 中间张量的口径同步改成全量。一句话：边界还上游一个全量，stage 内换回真正的 FlashComm1。改后的样子见图 4：

4. 主线讲解：从架构全貌到这次 PR

先对齐地基。本章走两条线：线一·模型架构线——从模型架构看到 vllm-ascend 框架，看模型的每一件住哪、谁调它，4.1 先摆这块推理的朴素 PyTorch 写法、定位到框架加的壳（图 5 给出全局定位 + 4.2 目录地图·图 6）；线二·PyTorch 实现线——从十行示例的 PyTorch 实现看到框架实现，看它被谁一步步算出来、每步为什么（4.3 流水线·图 7 → 4.4 全链时序·图 8 → 4.5 逐步收益·图 9）；两条线在 4.6 的问题处交汇，4.7 讲这次 PR（图 10~14），4.8~4.10 三张表收拢，每张图配一段讲解。先说结论：attention 的数学一行没动；本 PR 改的全是"开关与边界"——层的 SP 开关、o_proj 的归约权、PP 出口的布局。

4.1 先看这块推理的原始写法：朴素 PyTorch

老规矩，先立对比基准——这块推理（一层 Transformer 的前向）用 PyTorch 朴素写出来长什么样。自己动手写模型推理就会这么写：单卡、没有任何并行概念（有删节）。十行示例管的是 attention 这一小块；这里把整层摆出来，MoE 一并看：

# 朴素 PyTorch 写法：一层 Transformer 的前向（单卡，无框架）
def layer_forward(hidden_states):
    h = input_layernorm(hidden_states)          # 进层先归一化
    h = self_attn(h)                            # 注意力：QKV 投影 → QKᵀ/softmax/·V → o_proj（十行示例就是这段）
    hidden_states = h + hidden_states           # 残差 1
    h = post_attention_layernorm(hidden_states) # MoE 前再归一化
    h = mlp(hidden_states=h)                    # 前馈；GLM 这层是 MoE：路由选专家 → 专家 FFN → 加权合并
    return h + hidden_states                    # 残差 2 → 交给下一层

# 模型级就是 for 循环把 N 层串起来，最后 norm + lm_head 出下一个 token

这次 PR 动的不是这段数学——数学一行没动。动的是 vLLM 框架给这段朴素代码加的壳：进了框架，这一层被接上三件并行零件——层的 SP 开关（Ⓑ）、o_proj 的归约接线（Ⓐ）、PP 出口的交接。裸 PyTorch 里根本没有这些东西，它们全是框架为多卡布局塞进来的；PP>1 时框架还一刀切关掉 SP，让 Ⓐ 每层白付一次归约，而 FlashComm1 配置位照样报开启——名实不符，问题全住在壳上（ⒶⒷ 编号见第 2 章）。

这段代码在全局的什么位置？把模型架构到框架的全貌画出来（图 5）：

看懂这张图，再定位本 PR。架构件一个没变、链子一段没短——变的是"运行时补丁 patch/" 这一层的两个开关：层外壳的 SP 开关（use_sequence_parallel_moe）和 o_proj 的归约开关（reduce_results）。这就是为什么图 5 的 ★ 落在层外壳行，而不是任何注意力算子上。

4.2 先认目录：与 attention 相关的目录，一个一个认

四个目录，一条链，外加一个补丁层。worker/ 管执行与喂料，models/ 管模型结构，ops/ 管算子门面，attention/ 管注意力实现，patch/ 管运行时补丁——本 PR 的两个源码文件恰好一个在 patch/、一个在 worker/：它不改结构、不改算子，改的是"结构与算子被怎么组装、边界按什么口径喂"。

4.3 实现线：attention 在 vllm-ascend 原本是怎么完成的

架构线走完，现在切到实现线：attention 数学（图 7 左列的十行，L1~L10）在 vllm-ascend 里被一步步完成——左列每行，中列都有落点：

看懂这张图，记住三件事。

一、十行示例一步没少。左列每一行，中列都有对应——投影、RoPE、写缓存、稀疏注意力、收口，数学原样；唯一的例外是 L5（GQA repeat），MLA 系用 latent KV 全头共享达到同样目的，所以没有独立步骤。

二、实现分散在三个模块，但只有一条调用链。models/ 发起 → ops/ 放行 → attention/ 干活；右列就是每一站的地盘和门牌（文件:行号）。

三、本 PR 不碰这张图。九步流水线是"stage 内怎么算"，本 PR 管的是"层外壳的开关"与"PP 边界的布局"——往下看时序图时请带着这个分工。

4.4 全链时序：谁调用谁，每一棒先自己干再交棒

图 7 看的是"数学落在哪"。这份工作是谁发起、又一层层交给谁的？把整条调用链按时间拉开——六个类六条泳道，每个函数先标"自己干了啥"，再标"委托给谁"：

读这张图抓三个节奏。

一、委托链只有一条，没有岔路。引擎组批 → Worker 接批 → Runner 喂料 → 层外壳发起 → 门卫放行 → 后端开算；每一棒只把活交给下一棒，谁也不越级。

二、"自己干"的都是壳上的杂活，"委托"出去的才是正餐。Worker 自己只做监控与流水线收发；Runner 自己只做 pad 和组元数据；把十行示例真正算出来的，是最右边后端的那串自转圈。

三、本 PR 动的是泳道之外的两样东西。层外壳进来的姿势（SP 开没开、o_proj 归不归约）和 PP 出口交出去的姿势（全量还是分片）——泳道内的每一棒都没变。

4.5 为什么这么实现：逐步收益

图 7 回答了"落在哪"，图 8 回答了"谁调用谁"，还剩最后一个问题：每一步为什么这么实现？各自省的是什么？把图 7 中列的流水线原样搬到最左列，向右逐层展开——怎么做（带门牌）→ 为什么这么设计 → 收益点：

逐行看收益，"省"分三种，省的东西不一样。

省显存的：低秩 latent 让 KV Cache 不随头数翻倍；slot 页表定位省掉整块 cat 的大拷贝。

省带宽的：融合算子让中间结果不落显存——投影+rmsnorm+rope+写缓存一个算子，QKᵀ→softmax→·V 一个算子。

省算力的：稀疏注意力只算 indexer 选中的 top-2048 块，不扫全量历史；序列越长，省得越多。

而本 PR 要省的，是这张图之外的一次通信：FlashComm1 语义下 o_proj 本可以不归约（每层省一次集合通信）——PP>1 时被 guard 一刀切关掉。九步数学没变，变的是开关。

4.6 问题：PP>1 时，SP 被 guard 关掉了

两条线在这里交汇：架构线上是 Ⓑ 开关问题，落到实现线上，它打在层外壳与 PP 边界——十行示例本身无恙。把问题钉在代码上。上游的 guard 长在层外壳的组装处：PP>1 时 use_sequence_parallel_moe 不会被置真，decoder SP 整段关闭；而 o_proj 是 RowParallelLinear，reduce_results=True 意味着每层投影后各自做一次归约——这在与 FlashComm1 语义（归约收归 decoder 统一做）组合时就是纯白做：

# 改前状态（上游 guard 生效时，有删节）
layer.use_sequence_parallel_moe = False        # Ⓑ guard：PP>1 时 decoder SP 被关
layer.self_attn.o_proj.reduce_results = True   # Ⓐ 每层投影后各归各的——白付一次集合通信
# 而配置侧 FlashComm1 仍是"开启"——名实不符就从这来

更隐蔽的是这份"名实不符"的迷惑性：看配置、看日志，FlashComm1 都像是在跑；只有对照层外壳才知道 SP 没开、归约白付。问题不在十行数学，在开关与它背后的前提。

4.7 这次 PR：五条件重开 SP，边界还上游一个全量

PR 的三处源码改动，一句话一个：

# PR 新增（patch_deepseek_v2.py :326，有删节）：模型 init 补丁里的选路函数
def _enable_glm_pp_sequence_parallel(self, vllm_config) -> None:
    # 上游在 PP 下关 decoder SP，因为 transport 期望全量——
    # 本补丁的 forward 会在每个 PP 出口恢复全量，所以这里敢开
    if vllm_config.use_v2_model_runner:
        return                                     # 条件1：仅 V1 runner
    parallel_config = vllm_config.parallel_config
    if (vllm_config.model_config.hf_text_config.model_type != "glm_moe_dsa"
            or getattr(vllm_config.model_config.hf_text_config, "moe_layer_freq", 1) != 1
            or parallel_config.pipeline_parallel_size <= 1
            or not parallel_config.use_sequence_parallel_moe):
        return                                     # 条件2~5：GLM DSA / 层频=1 / PP>1 / SP 开
    for layer in islice(self.layers, self.start_layer, self.end_layer):
        if isinstance(layer.mlp, DeepseekV2MoE) and isinstance(layer.self_attn, DeepseekV2MLAAttention):
            layer.use_sequence_parallel_moe = True       # 重开 stage 内 SP（dense 层不动）
            layer.self_attn.o_proj.reduce_results = False  # 交出归约权：同一投影对象被 MLA 后端持有

选路是五条件与门（图 13）：任一不满足就整段返回，行为与改前完全一致。动到的层必须"mlp 是 MoE 且 self_attn 是 MLA"——dense 前缀层天然被跳过。

这张图就是整个 PR。左列是原本（PP>1 的跑法），中列是 PR 的做法，逐行对着看：根源行（guard）由五条件与门取代；代价行（Ⓐ 白归约）由归约权上交消掉；矛盾行（SP 开则出口分片）由聚全量上移化解；口径行（Runner）由新 helper 统一；最后两行是 scope 与闸门——不满足条件就回老路。

第二处源码改动是聚散搬移：同一段"聚回全量"的代码（cat → all_gather → 截断 → split），从"最后一个 rank 的 norm 之前"上移到"每个 PP 出口的 transport 之前"——改前 SP 关着这段代码形同虚设，改后 SP 开了它就是每个出口的必经站：

4.8 对照总表：原本 vs 这次 PR，逐项对照

4.9 PR 落点总表

4.10 测试（3 个文件，选路矩阵 + 口径匹配两道锁）

5. 前后代码对照（带注释）

"改后"代码均摘自 PR 实际 diff（有删节，只保留主干）；行号以 PR diff 为准。绿色为本 PR 新增逻辑；ⒶⒷ 编号对应图 2 / 图 4。

5.1 模型 init 补丁：五条件重开 SP、交归约权

def _patched_deepseek_v2_model_init(self, *args, **kwargs):
    _original_deepseek_v2_model_init(self, *args, **kwargs)
    # （本 PR 之前）这里没有任何 PP/SP 选路——
    # PP>1 时层 SP 维持上游 guard 的关闭状态，Ⓑ 生效
    self._use_upstream_aux_relay = ...
    self.topk_indices_buffer = ...

def _patched_deepseek_v2_model_init(self, *args, **kwargs):
    _original_deepseek_v2_model_init(self, *args, **kwargs)
    _enable_glm_pp_sequence_parallel(self, kwargs["vllm_config"])  # ← 新增调用
    ...

def _enable_glm_pp_sequence_parallel(self, vllm_config) -> None:
    if vllm_config.use_v2_model_runner:
        return                                    # 条件1：仅 V1 runner
    parallel_config = vllm_config.parallel_config
    if (vllm_config.model_config.hf_text_config.model_type != "glm_moe_dsa"
            or getattr(vllm_config.model_config.hf_text_config, "moe_layer_freq", 1) != 1
            or parallel_config.pipeline_parallel_size <= 1
            or not parallel_config.use_sequence_parallel_moe):
        return                                    # 条件2~5
    for layer in islice(self.layers, self.start_layer, self.end_layer):
        if isinstance(layer.mlp, DeepseekV2MoE) and isinstance(layer.self_attn, DeepseekV2MLAAttention):
            layer.use_sequence_parallel_moe = True       # 重开 SP（dense 层被 isinstance 挡住）
            layer.self_attn.o_proj.reduce_results = False  # 归约权上交；同一投影对象被 Ascend MLA 后端持有

5.2 PP 出口：聚全量上移到 transport 之前（图 14）

    if not pp_group.is_last_rank:
        pp_output = IntermediateTensors(...)
        ...
        return pp_output
    # ↓ 下面这段只在最后一个 rank 执行得到
    if hidden_states.shape[0] != positions.shape[0]:
        combined = torch.cat([hidden_states, residual], dim=-1)
        combined = tensor_model_parallel_all_gather(combined, 0)
        ...
    hidden_states, _ = self.norm(hidden_states, residual)

    # Restore full tokens before PP transport as well as the final norm.
    if hidden_states.shape[0] != positions.shape[0]:
        combined_states = torch.cat([hidden_states, residual], dim=-1)
        combined_states = tensor_model_parallel_all_gather(combined_states, 0)
        combined_states = combined_states[: positions.shape[0]]
        hidden_size = self.hidden_size
        hidden_states, residual = combined_states.split([hidden_size, hidden_size], dim=-1)
        residual = residual.contiguous()
    if not pp_group.is_last_rank:
        pp_output = IntermediateTensors({"hidden_states": hidden_states, "residual": residual})
        ...                                      # 出口交出去的必是全量
    hidden_states, _ = self.norm(hidden_states, residual)

5.3 Runner：PP 中间张量口径统一（图 15）

def _pp_intermediate_token_count(self, num_tokens: int) -> int:
    # GLM 的 decoder 在 stage 内用 SP，但补丁会在每个 PP 边界恢复全量
    if self.model_config.hf_text_config.model_type == "glm_moe_dsa" or not enable_sp():
        return num_tokens                        # GLM：全量
    tp = self.vllm_config.parallel_config.tensor_parallel_size
    return (num_tokens + tp - 1) // tp           # 其它 SP 模型：分片 ceil(n/tp)
# 三处硬编码替换：sync_and_slice_intermediate_tensors 的接收切片、
# _dummy_run 的显存预估与最大 buffer——口径从此只有一份

5.4 归约权转移：一行改动的分量

layer.self_attn.o_proj.reduce_results = False
# 别看只是一行:o_proj 是 RowParallelLinear,reduce_results=True 时
# 每层投影后各自归约(Ⓐ);置 False 后归约权交给 decoder 统一收口——
# 注释原话:This is the same projection object held by the Ascend MLA backend.
# 即改的这个对象,就是门卫/后端链上真正干活的那个投影——改一处,全链生效

6. 与我们排查问题的关系

方法论上值得记一笔：guard 的约束（transport 期望全量）是对的，错的是一刀切的实现。本 PR 的解法不是绕过约束，而是把约束的前提自己兑现——"边界还全量，stage 内才敢开 SP"。这个模式和 18671 的"布局契约"一脉相承：与其让开关一刀切，不如把布局责任钉在边界上。

PR 自述状态：ruff 通过；选路矩阵与口径匹配的回归用例已就位；张量单测与 NPU 正确性/性能验证仍 pending，此前一轮 CI 败在环境（pip 证书与 apt 签名），还没跑到测试。

7. 收益、风险与待验证

8. 术语速查

来源：github.com/vllm-project/vllm-ascend PR #18870（开放中，base: releases/v0.30.0，+176/−28，5 文件）· 本页为技术解读，行号以 PR diff 为准，数据均引自 PR 原文